# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/4 10:37 
@Author  : ZhangShenao 
@File    : 8.利用断点实现Agent的人机交互.py 
@Desc    : 利用断点实现Agent的人机交互
人在环路(Human-in-the-loop,简称 HIL)交互对于Agent系统至关重要
特别是在一些特定领域的Agent中,需要经过人类的允许或者指示才能进入下一步(例如某些敏感或者重要操作)
HIL最重要的部分就是interrupt断点

Interrupt建立在LangGraph的Checkpoint之上
Checkpoint在每个节点执行后保存图的状态,并且可以在指定的interrupt处暂停,等待人为批准
带收到人类确认后,从最后一个Checkpoint恢复执行

在LangGraph中可以通过在compile()函数
在编译的时候传递interrupt_before(前置断点)或者interrupt_after(后置断点)
这样在图结构程序执行到特定的断点时就会暂停执行,等待其他操作(例如人类确认，修改状态等)

如果需要恢复图执行,只需要再次调用invoke/stream等,并传递inputs=None
传递输入为None意味着像中断没有发生一样继续执行
基于这个思路就可以实现让人类干预图的执行
"""
from typing import Literal

import dotenv
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import END
from langgraph.graph import MessagesState, StateGraph
from langgraph.prebuilt import ToolNode

# 加载环境变量
dotenv.load_dotenv()

# 构造工具列表
import os

import dotenv
from langchain_openai import ChatOpenAI

from dalle_image_tool import create_dalle_image_tool
from gaode_weather_tool import GaodeWeatherTool
from google_serper_tool import create_google_serper_tool

# 加载环境变量
dotenv.load_dotenv()

# 创建工具列表
tools = [GaodeWeatherTool(), create_dalle_image_tool(), create_google_serper_tool()]

# 创建LLM,并绑定工具列表
llm = (ChatOpenAI(model="gpt-4o",
                  temperature=0,
                  openai_api_base=os.getenv("OPENAI_API_BASE"))
       .bind_tools(tools=tools))


# 创建LLM节点
def llm_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
    inputs = state["messages"]
    replies = llm.invoke(inputs)
    return {"messages": [replies]}


# 创建条件路由函数
def condition_route(state: MessagesState, config: RunnableConfig) -> Literal["tools", "__end__"]:
    # 获取最后一条AIMessage
    ai_msg = state["messages"][-1]

    # 如果AI消息中包含工具调用信息,则跳转到tools节点,执行工具调用
    if hasattr(ai_msg, "tool_calls") and len(ai_msg.tool_calls) > 0:
        return "tools"

    # 否则为普通消息,说明工具调用已经完成,直接跳转到END节点
    return END


# 初始化Graph,指定状态为MessagesState
graph_builder = StateGraph(state_schema=MessagesState)

# 添加节点
graph_builder.add_node(node="llm", action=llm_node)  # 添加LLM节点
# 添加工具节点
# 采用LangGraph预构建的ToolNode节点,并指定工具列表
graph_builder.add_node(node="tools", action=ToolNode(tools=tools))

# 添加边
# 设置起始节点
graph_builder.set_entry_point(key="llm")
# 设置条件边,并指定条件路由函数
graph_builder.add_conditional_edges(source="llm", path=condition_route)
# 设置普通边
graph_builder.add_edge(start_key="tools", end_key="llm")

# 创建CheckPointer组件,用于持久化Graph的状态
# 将状态数据保存在内存中
checkpointer = MemorySaver()

# 将Graph编译为Runnable可运行组件,并指定checkpointer
# 设置interrupt_before=tools,即在执行工具调用节点前中断,等待人工确认
graph = graph_builder.compile(checkpointer=checkpointer,
                              interrupt_before=["tools"])

# 定义运行时配置,并指定Checkpointer绑定的thread_id
config = {"configurable": {"thread_id": 1}}

# 定义初始状态
init_state = {"messages": [
    ("human", "2024年奥运会男篮前3名分别是哪个国家？")
]}

# 执行Graph,并传入运行时配置
# 由于设置了interrupt,因此在执行tools节点前,Graph会暂停,等待人类确认
middle_state = graph.invoke(input=init_state, config=config)
print(f"执行到Graph断点, 中间状态:\n{middle_state}")

# 等待人类确认
tool_calls = middle_state["messages"][-1].tool_calls
print(f"需要调用工具: {tool_calls}")
confirm = input("是否确认调用？(yes/no) ")

if confirm.lower() == "yes":
    print("确认调用工具")
    # 确认调用工具
    # 传递input=None,Graph会从中断中恢复
    # 如果传递input为其它状态,则会对状态进行更新
    # 利用CheckPointer恢复Graph状态,并继续执行
    final_state = graph.invoke(input=None, config=config)
    ai_reply = final_state["messages"][-1].content
    print(f"Graph执行结束，最终结果:\n{ai_reply}")
else:
    print("取消调用工具,Graph执行结束")
